package com.hugy.test.entity.es;

import com.google.common.collect.Lists;
import lombok.Data;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramInterval;
import org.joda.time.DateTimeZone;

import java.util.Arrays;
import java.util.List;

/**
 * @author YangChao
 * @create 2019-03-27 11:03
 **/
@Data
public class AggregationParam extends QueryParam {

    private List<AggregationBuilder> groups = Lists.newArrayList();

    private List<AggregationBuilder> aggregations = Lists.newArrayList();

    public static AggregationParam empty() {
        return new AggregationParam();
    }

    public void  groupBy(String... fields) {
        Arrays.stream(fields).map(field -> AggregationBuilders.terms(field).field(field).size(Integer.MAX_VALUE)).forEach(this.groups::add);
    }

    public void groupByAs(String field, String alias) {
        this.groups.add(AggregationBuilders.terms(alias).field(field).size(Integer.MAX_VALUE));
    }

    public void groupByAs(String field, String alias, Integer limit) {
        this.groups.add(AggregationBuilders.terms(alias).field(field).size(limit == null ? Integer.MAX_VALUE : limit));
    }

    public void groupBy(DateHistogramInterval interval, String... fields) {
        Arrays.stream(fields).map(field -> AggregationBuilders.dateHistogram(field).dateHistogramInterval(interval).field(field).minDocCount(1).timeZone(DateTimeZone.forOffsetHours(8))).forEach(this.groups::add);
    }

    public void groupBy(DateHistogramInterval interval, String formatter, String field) {
        this.groups.add(AggregationBuilders.dateHistogram(field).dateHistogramInterval(interval).field(field).minDocCount(1).format(formatter).timeZone(DateTimeZone.forOffsetHours(8)));
    }

    public void groupByAs(DateHistogramInterval interval, String field, String alias) {
        this.groups.add(AggregationBuilders.dateHistogram(alias).dateHistogramInterval(interval).field(field).minDocCount(1).timeZone(DateTimeZone.forOffsetHours(8)));
    }

    public void groupByAs(DateHistogramInterval interval, String formatter, String field, String alias) {
        this.groups.add(AggregationBuilders.dateHistogram(alias).dateHistogramInterval(interval).field(field).minDocCount(1).format(formatter).timeZone(DateTimeZone.forOffsetHours(8)));
    }

    public void groupByAs(Integer precision, String field) {
        groupByAs(precision, field, field);
    }

    public void groupByAs(Integer precision, String field, String alias) {
        this.groups.add(AggregationBuilders.geohashGrid(alias).precision(precision).field(field).size(Integer.MAX_VALUE));
    }

    public void groupBy(AggregationBuilder... fields) {
        this.groups.addAll(Arrays.asList(fields));
    }

    public void aggregation(AggregationField... aggregationFields) {
        Arrays.stream(aggregationFields).map(a -> {
            String field = a.getField();
            String alias = a.getAlias();
            switch (a.getType()) {
                case AVG:
                    return AggregationBuilders.avg(alias).field(field);
                case MAX:
                    return AggregationBuilders.max(alias).field(field);
                case MIN:
                    return AggregationBuilders.min(alias).field(field);
                case SUM:
                    return AggregationBuilders.sum(alias).field(field);
                case COUNT:
                    return AggregationBuilders.count(alias).field(field);
                default:
                    return AggregationBuilders.count(alias).field(field);
            }
        }).forEach(this.aggregations::add);
    }

    public void aggregation(AggregationBuilder... aggregations) {
        this.aggregations.addAll(Arrays.asList(aggregations));
    }

}